{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a6fce64",
   "metadata": {},
   "source": [
    "# 合法化 qdq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0c1d3b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入必要的模块和类\n",
    "import tvm\n",
    "from tvm.relax.transform import LegalizeOps\n",
    "from tvm.script import relax as R, tir as T\n",
    "import tvm.testing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f33171",
   "metadata": {},
   "source": [
    "## 测试从 float32 到 int8 的量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "db5e7abf",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Quantize:\n",
    "    @R.function\n",
    "    def main(\n",
    "        data: R.Tensor((2, 4), \"float32\"),\n",
    "        scale: R.Tensor((2,), \"float32\"),\n",
    "        zp: R.Tensor((2,), \"int8\"),\n",
    "    ) -> R.Tensor((2, 4), \"int8\"):\n",
    "        # 调用量化操作，沿轴0量化到int8\n",
    "        out = R.quantize(data, scale, zp, axis=0, out_dtype=\"int8\")\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a77e9e45",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">quantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), scale: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), zp: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), quantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;quantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], scale[v_i0], zp[v_i0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(quantized[v_i0, v_i1])\n",
       "                quantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>round(data[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">/</span> scale[v_i0]) <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, zp[v_i0]), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">127.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">128.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), scale: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), zp: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>quantize, (data, scale, zp), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Quantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eddd477",
   "metadata": {},
   "source": [
    "## 测试从 float16 到 uint8 的量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d75a493c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Quantize:\n",
    "    @R.function\n",
    "    def main(\n",
    "        data: R.Tensor((2, 4), \"float16\"),\n",
    "        scale: R.Tensor((2,), \"float16\"),\n",
    "        zp: R.Tensor((2,), \"int8\"),\n",
    "    ) -> R.Tensor((2, 4), \"uint8\"):\n",
    "        # 调用量化操作，沿轴0量化到uint8\n",
    "        out = R.quantize(data, scale, zp, axis=0, out_dtype=\"uint8\")\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9a2b20a0",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">quantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float16&quot;</span>), scale: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;float16&quot;</span>), zp: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), quantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;uint8&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;quantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], scale[v_i0], zp[v_i0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(quantized[v_i0, v_i1])\n",
       "                quantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;uint8&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>round(data[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">/</span> scale[v_i0]) <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float16&quot;</span>, zp[v_i0]), T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #008000\">255.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #008000\">0.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>), scale: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>), zp: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;uint8&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>quantize, (data, scale, zp), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;uint8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Quantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56087eb8",
   "metadata": {},
   "source": [
    "## 测试符号形状输入下，从float32到int8的量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cf19d467",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Quantize:\n",
    "    @R.function\n",
    "    def main(\n",
    "        data: R.Tensor((4, \"n\"), \"float32\"),\n",
    "        scale: R.Tensor((\"n\",), \"float32\"),\n",
    "        zp: R.Tensor((\"n\",), \"int8\"),\n",
    "    ) -> R.Tensor((4, \"n\"), \"int8\"):\n",
    "        # 调用量化操作，沿最后一个轴(-1)量化到int8\n",
    "        out = R.quantize(data, scale, zp, axis=-1, out_dtype=\"int8\")\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f1fdd0ca",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">quantize</span>(var_data: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, var_scale: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, var_zp: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, var_quantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        data <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_data, (T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>), n))\n",
       "        scale <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_scale, (n,))\n",
       "        zp <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_zp, (n,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "        quantized <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_quantized, (T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>), n), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>), n):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;quantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], scale[v_i1], zp[v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(quantized[v_i0, v_i1])\n",
       "                quantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>round(data[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">/</span> scale[v_i1]) <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, zp[v_i1]), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">127.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">128.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">4</span>, <span style=\"color: #BA2121\">&quot;n&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), scale: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), zp: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">4</span>, <span style=\"color: #BA2121\">&quot;n&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>):\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>quantize, (data, scale, zp), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">4</span>, n), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Quantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e4a79b9",
   "metadata": {},
   "source": [
    "## 测试使用标量参数（而非张量）时，从float32到int8的量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f58c798",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Quantize:\n",
    "    @R.function\n",
    "    def main(data: R.Tensor((2, 4), \"float32\")) -> R.Tensor((2, 4), \"int8\"):\n",
    "        # 使用标量值作为scale和zp参数\n",
    "        out = R.quantize(\n",
    "            data, R.const(2.0, \"float32\"), R.const(1, \"int8\"), axis=-1, out_dtype=\"int8\"\n",
    "        )\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ba1db281",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">quantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), quantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;quantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(quantized[v_i0, v_i1])\n",
       "                quantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>round(data[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.5</span>)) <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1.0</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">127.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">128.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>quantize, (data,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Quantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16e8fbef",
   "metadata": {},
   "source": [
    "## 测试使用一维常量数组作为参数时，从float32到int8的量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "135efa87",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Quantize:\n",
    "    @R.function\n",
    "    def main(data: R.Tensor((2, 4), \"float32\")) -> R.Tensor((2, 4), \"int8\"):\n",
    "        # 使用一维常量数组作为scale和zp参数\n",
    "        out = R.quantize(\n",
    "            data,\n",
    "            R.const([2.0, 1.0], \"float32\"),\n",
    "            R.const([4, 5], \"int8\"),\n",
    "            axis=0,\n",
    "            out_dtype=\"int8\",\n",
    "        )\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "effe929b",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">quantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), quantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;quantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], B[v_i0], C[v_i0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(quantized[v_i0, v_i1])\n",
       "                quantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>round(data[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">/</span> B[v_i0]) <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, C[v_i0]), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">127.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">128.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>quantize, (data, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>]), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "\n",
       "<span style=\"color: #007979; font-style: italic\"># Metadata omitted. Use show_meta=True in script() method to show it.</span>\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Quantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a80abbb8",
   "metadata": {},
   "source": [
    "## 测试使用标量参数时，从float16到int8的量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3c31e80a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Quantize:\n",
    "    @R.function\n",
    "    def main(data: R.Tensor((2, 4), \"float16\")) -> R.Tensor((2, 4), \"int8\"):\n",
    "        # 使用标量值作为scale和zp参数\n",
    "        out = R.quantize(\n",
    "            data, R.const(2.0, \"float16\"), R.const(1, \"int8\"), axis=-1, out_dtype=\"int8\"\n",
    "        )\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8cfd370e",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">quantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float16&quot;</span>), quantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;quantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(quantized[v_i0, v_i1])\n",
       "                quantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>round(data[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #008000\">0.5</span>)) <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #008000\">1.0</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #008000\">127.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">128.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>quantize, (data,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Quantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a8483b8",
   "metadata": {},
   "source": [
    "## 测试从int8到float32的反量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f400138b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Dequantize:\n",
    "    @R.function\n",
    "    def main(\n",
    "        data: R.Tensor((2, 4), \"int8\"),\n",
    "        scale: R.Tensor((2,), \"float32\"),\n",
    "        zp: R.Tensor((2,), \"int8\"),\n",
    "    ) -> R.Tensor((2, 4), \"float32\"):\n",
    "        # 调用反量化算子，沿轴0反量化到float32\n",
    "        out = R.dequantize(data, scale, zp, axis=0, out_dtype=\"float32\")\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7949a863",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">dequantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), scale: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), zp: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), dequantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;dequantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], zp[v_i0], scale[v_i0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(dequantized[v_i0, v_i1])\n",
       "                dequantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data[v_i0, v_i1]) <span style=\"color: #A2F; font-weight: bold\">-</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, zp[v_i0])) <span style=\"color: #A2F; font-weight: bold\">*</span> scale[v_i0]\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>), scale: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), zp: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>dequantize, (data, scale, zp), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Dequantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9831159b",
   "metadata": {},
   "source": [
    "## 测试使用标量参数时，从int8到float32的反量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d83d19f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Dequantize:\n",
    "    @R.function\n",
    "    def main(data: R.Tensor((2, 4), \"int8\")) -> R.Tensor((2, 4), \"float32\"):\n",
    "        # 使用标量值作为scale和zp参数\n",
    "        out = R.dequantize(\n",
    "            data, R.const(2.0, \"float32\"), R.const(1, \"int8\"), axis=0, out_dtype=\"float32\"\n",
    "        )\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8319ef13",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">dequantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), dequantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;dequantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(dequantized[v_i0, v_i1])\n",
       "                dequantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data[v_i0, v_i1]) <span style=\"color: #A2F; font-weight: bold\">-</span> <span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2.0</span>)\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>dequantize, (data,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Dequantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9c7fc20",
   "metadata": {},
   "source": [
    "## 测试符号形状输入下，从int8到float32的反量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c0ce2bf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Dequantize:\n",
    "    @R.function\n",
    "    def main(\n",
    "        data: R.Tensor((2, \"n\"), \"int8\"),\n",
    "        scale: R.Tensor((\"n\",), \"float32\"),\n",
    "        zp: R.Tensor((\"n\",), \"int8\"),\n",
    "    ) -> R.Tensor((2, \"n\"), \"float32\"):\n",
    "        # 调用反量化操作，沿最后一个轴(-1)反量化到float32\n",
    "        out = R.dequantize(data, scale, zp, axis=-1, out_dtype=\"float32\")\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6c9467f4",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">dequantize</span>(var_data: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, var_scale: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, var_zp: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, var_dequantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        data <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_data, (T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), n), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "        scale <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_scale, (n,))\n",
       "        zp <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_zp, (n,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "        dequantized <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>match_buffer(var_dequantized, (T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), n))\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), n):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;dequantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], zp[v_i1], scale[v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(dequantized[v_i0, v_i1])\n",
       "                dequantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data[v_i0, v_i1]) <span style=\"color: #A2F; font-weight: bold\">-</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, zp[v_i1])) <span style=\"color: #A2F; font-weight: bold\">*</span> scale[v_i1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;n&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>), scale: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), zp: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;n&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>dequantize, (data, scale, zp), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, n), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Dequantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ab4508a",
   "metadata": {},
   "source": [
    "## 测试从int8到float16的反量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3cbe1d4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Dequantize:\n",
    "    @R.function\n",
    "    def main(\n",
    "        data: R.Tensor((2, 4), \"int8\"),\n",
    "        scale: R.Tensor((2,), \"float16\"),\n",
    "        zp: R.Tensor((2,), \"int8\"),\n",
    "    ) -> R.Tensor((2, 4), \"float16\"):\n",
    "        # 调用反量化操作，沿轴0反量化到float16\n",
    "        out = R.dequantize(data, scale, zp, axis=0, out_dtype=\"float16\")\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9a6a4513",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">dequantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), scale: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;float16&quot;</span>), zp: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>),), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), dequantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float16&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;dequantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1], zp[v_i0], scale[v_i0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(dequantized[v_i0, v_i1])\n",
       "                dequantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float16&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data[v_i0, v_i1]) <span style=\"color: #A2F; font-weight: bold\">-</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, zp[v_i0])) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, scale[v_i0]), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">65504.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">65504.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>), scale: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>), zp: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>dequantize, (data, scale, zp), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Dequantize)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "454e168f",
   "metadata": {},
   "source": [
    "## 测试使用标量参数时，从int8到float16的反量化算子的合法化变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "51d94c83",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Dequantize:\n",
    "    @R.function\n",
    "    def main(data: R.Tensor((2, 4), \"int8\")) -> R.Tensor((2, 4), \"float16\"):\n",
    "        # 使用标量值作为scale和zp参数\n",
    "        out = R.dequantize(\n",
    "            data, R.const(2.0, \"float16\"), R.const(1, \"int8\"), axis=0, out_dtype=\"float16\"\n",
    "        )\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b886c905",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">dequantize</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), dequantized: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)), <span style=\"color: #BA2121\">&quot;float16&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;dequantized&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(data[v_i0, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(dequantized[v_i0, v_i1])\n",
       "                dequantized[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float16&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T<span style=\"color: #A2F; font-weight: bold\">.</span>min(T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;float32&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data[v_i0, v_i1]) <span style=\"color: #A2F; font-weight: bold\">-</span> <span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2.0</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">65504.0</span>)), T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">65504.0</span>)))\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        out <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>dequantize, (data,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换并验证结果是否符合预期\n",
    "mod = LegalizeOps()(Dequantize)\n",
    "mod.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
